package chapter5.eg4;

import chapter5.AbstractConsumer;
import java.time.Duration;
import java.util.Collection;
import java.util.Collections;
import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.ArrayBlockingQueue;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.ThreadPoolExecutor;
import java.util.concurrent.ThreadPoolExecutor.CallerRunsPolicy;
import java.util.concurrent.TimeUnit;
import org.apache.kafka.clients.consumer.ConsumerConfig;
import org.apache.kafka.clients.consumer.ConsumerRebalanceListener;
import org.apache.kafka.clients.consumer.ConsumerRecords;
import org.apache.kafka.clients.consumer.KafkaConsumer;
import org.apache.kafka.clients.consumer.OffsetAndMetadata;
import org.apache.kafka.common.TopicPartition;
import org.apache.kafka.common.errors.WakeupException;

public class ConsumerThreadHandler<K, V> extends AbstractConsumer {

  private final KafkaConsumer<K, V> consumer;
  private ExecutorService executor;
  private final Map<TopicPartition, OffsetAndMetadata> offsets = new HashMap<>();

  public ConsumerThreadHandler(String groupId, String topic) {
    properties.put(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG, "false");
    consumer = new KafkaConsumer<>(properties);
    consumer.subscribe(Collections.singletonList(topic), new ConsumerRebalanceListener() {
      @Override
      public void onPartitionsRevoked(Collection<TopicPartition> partitions) {
        consumer.commitSync(offsets);
      }

      @Override
      public void onPartitionsAssigned(Collection<TopicPartition> partitions) {
        offsets.clear();
      }
    });
  }

  public void connsume(int threadNum) {
    executor = new ThreadPoolExecutor(threadNum, threadNum, 0L, TimeUnit.MILLISECONDS,
      new ArrayBlockingQueue<>(1000), new CallerRunsPolicy());
    try {
      while (true) {
        ConsumerRecords<K, V> records = consumer.poll(Duration.ofMillis(1000));
        if (!records.isEmpty()) {
          executor.submit(new ConsumerWorker<>(records, offsets));
        }
        commitOffsets();
      }
    } catch (WakeupException e) {
      // ignore
    } finally {
      commitOffsets();
      consumer.close();
    }
  }

  private void commitOffsets() {
    Map<TopicPartition, OffsetAndMetadata> unmodifiedMap;
    synchronized (offsets) {
      if (offsets.isEmpty()) {
        return;
      }

      unmodifiedMap = Collections.unmodifiableMap(new HashMap<>(offsets));

      offsets.clear();
    }

    consumer.commitSync(unmodifiedMap);
  }

  public void close() {
    consumer.wakeup();
    executor.shutdown();
  }

}
